from langchain.agents import AgentType, initialize_agent, Tool
from langchain_openai import OpenAI
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from langchain.memory import ConversationBufferMemory
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from datetime import datetime
import os

# 设置OpenAI API密钥
os.environ["OPENAI_API_KEY"] = "your-openai-api-key-here"

# 创建搜索工具
search = DuckDuckGoSearchAPIWrapper()

def duckduckgo_search(query):
    """使用DuckDuckGo搜索网络信息"""
    return search.run(query)

# 文档生成模板
document_template = """
根据以下搜索信息，生成一份结构化的文档：

搜索主题: {topic}
搜索时间: {timestamp}

搜索到的信息:
{search_results}

请生成一份包含以下内容的详细文档：
1. 执行摘要
2. 主要发现
3. 关键信息点
4. 相关链接（如果适用）
5. 结论和建议

文档内容："""

def generate_document(topic, search_results):
    """生成结构化文档"""
    prompt = PromptTemplate(
        template=document_template,
        input_variables=["topic", "timestamp", "search_results"]
    )
    
    llm = OpenAI(temperature=0.7)
    chain = LLMChain(llm=llm, prompt=prompt)
    
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    return chain.run({
        "topic": topic,
        "timestamp": timestamp,
        "search_results": search_results
    })

# 创建工具列表
tools = [
    Tool(
        name="Web Search",
        func=duckduckgo_search,
        description="用于搜索网络上的最新信息，当需要获取实时数据或事实信息时使用"
    ),
    Tool(
        name="Document Generator",
        func=generate_document,
        description="根据搜索到的信息生成结构化文档，输入格式为：主题,搜索结果"
    )
]

# 初始化LLM
llm = OpenAI(temperature=0)

# 创建记忆
memory = ConversationBufferMemory(memory_key="chat_history")

# 初始化agent
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
    memory=memory,
    verbose=True
)

def search_and_generate_document(topic):
    """完整的搜索和文档生成流程"""
    print(f"正在搜索主题: {topic}")
    
    # 第一步：搜索信息
    search_query = f"{topic} 最新信息 新闻 报道"
    search_results = duckduckgo_search(search_query)
    
    print("搜索完成，开始生成文档...")
    
    # 第二步：生成文档
    document = generate_document(topic, search_results)
    
    return document

# 使用示例
if __name__ == "__main__":
    # 测试完整的搜索和文档生成
    topic = "人工智能最新发展"
    
    try:
        document = search_and_generate_document(topic)
        print("\n" + "="*50)
        print("生成的文档:")
        print("="*50)
        print(document)
        
        # 保存文档到文件
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{topic.replace(' ', '_')}_{timestamp}.txt"
        
        with open(filename, 'w', encoding='utf-8') as f:
            f.write(document)
        
        print(f"\n文档已保存到: {filename}")
        
    except Exception as e:
        print(f"发生错误: {e}")
        print("请检查OpenAI API密钥是否正确设置")